#!usr/bin/env python  
# -*- coding:utf-8 -*-
""" 
@author:robot
@file: balance.py 
@version:
@time: 2024/01/25 
"""
"""
from top

基于二叉树的最大深度的基础，只要一棵二叉树的左右子树高度之差不大于1，并且保证该子树的所有子树都是平衡二叉树，即可判断其为平衡二叉树。
"""


class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


node1 = TreeNode(1)
node2 = TreeNode(2)
node3 = TreeNode(3)
node4 = TreeNode(4)
node5 = TreeNode(5)
node1.left = node2
node1.right = node3
node3.left = node4
node3.right = node5


def max_depth(root):
    if not root:
        return 0
    return max(max_depth(root.left), max_depth(root.right)) + 1


def balance_top(root):
    if not root:
        return True
    if abs(max_depth(root.left) - max_depth(root.right)) > 1:
        return False
    if not balance_top(root.left) or not balance_top(root.right):
        return False
    return True


print(balance_top(node1))
"""
from bottom

对于每一棵树，计算其左子树的高度与右子树的高度，当左右子树高度差不大于1时，向上层返回左右子树中最大高度值+1；
当发现任何一棵子树是不平衡树时，向上层返回-1.
"""


def balance_bottom(root):
    if dfs(root) == -1:
        return False
    return True


def dfs(root):
    if not root:
        return 0
    left_height = dfs(root.left)
    if left_height == -1:
        return -1
    right_height = dfs(root.right)
    if right_height == -1:
        return -1
    if abs(left_height - right_height) > 1:
        return -1
    return max(left_height, right_height) + 1


print(balance_bottom(node1))